package com.practice.niuke.new_direct_basics.class09;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/**
 * 并查集结构的详解和实现
 */
public class Code04_UnionFind {

    /**
     * 并查集结构中，每个集合的元素（节点）类
     *
     * @param <V>
     */
    public static class Element<V> {
        /**
         * 元素值
         */
        public V value;

        /**
         * 指向上一个节点的指针
         *
         * @param value
         */
        public Element(V value) {
            this.value = value;
        }

    }

    /**
     * 并查集结构
     *
     * @param <V>
     */
    public static class UnionFindSet<V> {
        // HashMap<节点a的值，节点a>：key为“节点a的值”，value为“节点a”
        public HashMap<V, Element<V>> elementMap;
        // HashMap<节点a, 节点a的父节点>
        public HashMap<Element<V>, Element<V>> fatherMap;
        // sizeMap中的key，每一个key都一定是集合的头节点（代表节点）。
        // HashMap<节点a, 节点a所代表的集合的元素个数>
        public HashMap<Element<V>, Integer> sizeMap;

        /**
         * 构造函数(初始化)
         *
         * @param list 样本集合（节点值集合）
         */
        public UnionFindSet(List<V> list) {
            elementMap = new HashMap<>();
            fatherMap = new HashMap<>();
            sizeMap = new HashMap<>();
            for (V value : list) {
                Element<V> element = new Element<V>(value);
                elementMap.put(value, element);
                fatherMap.put(element, element);
                sizeMap.put(element, 1);
            }
        }

        /**
         * 查询一个节点的头节点，
         * 从输入参数element出发，往上一直找，找到不能再往上的头节点返回。
         * 往上找的过程中，把经过的节点全部扁平化（都直接挂在头节点下）。
         *
         * @param element element
         * @return Element
         */
        private Element<V> findHead(Element<V> element) {
            // 把往上找的过程中，沿途经过的节点都记录在path栈中
            Stack<Element<V>> path = new Stack<>();
            // 一个节点的父节点不是自己就一直往上找，直到来到头节点（代表节点）。
            while (element != fatherMap.get(element)) {
                path.push(element);
                element = fatherMap.get(element);
            }
            // 把往上找的过程中，沿途经过的节点都直接挂在头节点（代表节点）下。
            while (!path.isEmpty()) {
                fatherMap.put(path.pop(), element);
            }
            return element;
        }

        /**
         * 判断两个元素是否属于同一个集合
         *
         * @param a 第一个元素的值
         * @param b 第二个元素的值
         * @return boolean
         */
        public boolean isSameSet(V a, V b) {
            if (elementMap.containsKey(a) && elementMap.containsKey(b)) {
                return findHead(elementMap.get(a)) == findHead(elementMap.get(b));
            }
            return false;
        }

        /**
         * 合并两个元素所在的集合
         *
         * @param a 第一个元素的值
         * @param b 第二个元素的值
         */
        public void union(V a, V b) {
            if (elementMap.containsKey(a) && elementMap.containsKey(b)) {
                Element<V> aF = findHead(elementMap.get(a));
                Element<V> bF = findHead(elementMap.get(b));
                if (aF != bF) {
                    Element<V> big = sizeMap.get(aF) >= sizeMap.get(bF) ? aF : bF;
                    Element<V> small = big == aF ? bF : aF;
                    fatherMap.put(small, big);
                    sizeMap.put(big, sizeMap.get(aF) + sizeMap.get(bF));
                    sizeMap.remove(small);
                }
            }
        }

    }

}
